# -*- coding: utf-8 -*-
# time: 2025/4/9 14:48
# file: data_ready.py
# author: hanson
"""
数据集准备
# 安装依赖
pip install langchain==0.3.0 torch==2.3.0 transformers==4.40.0 modelscope==1.13.0 peft==0.10.0
如果ImportError: cannot import name 'get_metadata_patterns' from 'datasets.data_files' 是 datasets版本过高
pip install datasets==3.1.0
数据集描述：https://www.modelscope.cn/datasets/AI-ModelScope/alpaca-gpt4-data-zh

"""

from modelscope import AutoTokenizer, AutoModelForCausalLM
from peft import LoraConfig, get_peft_model
import torch

# 1. 加载模型和tokenizer
#model_id = "qwen/qwen2-0.5b-instruct"
model_id = r"E:\soft\model\qwen\Qwen\Qwen2___5-0___5B-Instruct"
# AutoTokenizer.from_pretrained用于加载预训练的tokenizer，负责将文本转换为模型可理解的数字形式。
tokenizer = AutoTokenizer.from_pretrained(
    pretrained_model_name_or_path=model_id,
    #model_id,
                                          trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,   # model_id,
    torch_dtype=torch.bfloat16,  # 设置模型精度
    device_map="auto",  # 设置模型在哪个设备上运行，这里设置为自动选择
    trust_remote_code=True  # 是否信任远程代码，如果模型是私有的，需要设置为False
)

# 2. 添加LoRA适配器
peft_config = LoraConfig(
    r=8,    # r是LoRA算法中的超参数，它控制了权重矩阵的秩。
    lora_alpha=32,  # lora_alpha是LoRA算法中的超参数，它控制了权重矩阵的缩放因子。
    target_modules=["q_proj", "k_proj", "v_proj"],  # target_modules是LoRA算法中的关键参数，它指定了需要做LoRA适配的层。
    lora_dropout=0.05,      # lora_dropout是LoRA算法中的超参数，它控制了权重矩阵的随机化程度。
    bias="none",    # bias是LoRA算法中的超参数，它控制了权重矩阵的偏置项。
    task_type="CAUSAL_LM"   # task_type是LoRA算法中的超参数，它控制了权重矩阵的初始化方式。
)
model = get_peft_model(model, peft_config)

# 3. 准备训练数据（关键修改部分）
def tokenize_function(examples):
    # 构造训练格式：instruction + input -> output
    prompts = [
        f"Instruction: {inst}\nInput: {inp}\nOutput: {out}"
        for inst, inp, out in zip(
            examples["instruction"],   # instruction
            examples["input"],      # input
            examples["output"]
        )
    ]
    return tokenizer(prompts, truncation=True, max_length=512)

# 从ModelScope加载数据
from modelscope.msdatasets import MsDataset
dataset = MsDataset.load(
   # 'AI-ModelScope/alpaca-gpt4-data-zh',
    r'F:\temp\alpaca-gpt4-data-zh\train.csv',
    split='train'
).to_hf_dataset().select(range(4))  # 取1000条示例

tokenized_dataset = dataset.map(
    tokenize_function,  # 上面的tokenize_function 训练数据
    batched=True,
    remove_columns=dataset.column_names
)

# 4. 定义数据整理器（解决label_names问题）
from transformers import DataCollatorForLanguageModeling
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False  # 因果语言模型用mlm=False
)

# 5. 训练配置
from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./qwen2-0.5b-finetuned",
    per_device_train_batch_size=2,  #  4 设置batch_size 每个 GPU/CPU 的训练批次大小
    gradient_accumulation_steps=2,  # 设置梯度累积
    learning_rate=2e-5,  # 学习率
    num_train_epochs=1,  # 设置训练轮数 3
    logging_steps=10,  # 设置日志记录步数
    save_strategy="epoch",  # 设置保存策略
    fp16=False,  # 使用半精度 True  是否启用混合精度训练（需 GPU 支持）
    remove_unused_columns=False,
    label_names=["input_ids"]  # 明确指定标签列
)

trainer = Trainer(
    model=model,
    args=training_args,  # 训练参数
    train_dataset=tokenized_dataset,  ## 训练数据
    data_collator=data_collator  # 添加数据整理器
)

# 6. 开始训练
trainer.train()
model.save_pretrained("./qwen2-0.5b-finetuned")
tokenizer.save_pretrained("./qwen2-0.5b-finetuned")  # 必须显式保存tokenizer